package com.allwe.client.reBalance;

import lombok.Data;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;

import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 分区再均衡处理器
 *
 * @Author: AllWe
 * @Date: 2024/10/17/8:05
 */
@Data
public class ReBalanceHandler implements ConsumerRebalanceListener {

    // 记录每个分区的消费偏移量
    public final static ConcurrentHashMap<TopicPartition, Long> partitionOffsetMap = new ConcurrentHashMap<TopicPartition, Long>();

    private final Map<TopicPartition, OffsetAndMetadata> currOffsets;

    private final KafkaConsumer<String, String> consumer;

    public ReBalanceHandler(Map<TopicPartition, OffsetAndMetadata> currOffsets, KafkaConsumer<String, String> consumer) {
        this.currOffsets = currOffsets;
        this.consumer = consumer;
    }

    // 分区再均衡之前
    // 某一个消费者在让出分区之前，需要先将已消费的偏移量提交
    @Override
    public void onPartitionsRevoked(Collection<TopicPartition> collection) {
        // 线程id
        final String id = Thread.currentThread().getId() + "";
        System.out.println(id + "-onPartitionsRevoked参数值为：" + collection);
        System.out.println(id + "-服务器准备分区再均衡，提交偏移量。当前偏移量为:" + currOffsets);
        //我们可以不使用consumer.commitSync(currOffsets);
        //提交偏移量到kafka,由我们自己维护*/
        //开始事务
        //偏移量写入数据库
        System.out.println("分区偏移量表中:" + partitionOffsetMap);
        for (TopicPartition topicPartition : collection) {
            partitionOffsetMap.put(topicPartition, currOffsets.get(topicPartition).offset());
        }
        // 同步提交偏移量，等到成功后再往后执行
        consumer.commitSync(currOffsets);
    }

    // 分区再均衡之后
    // 新的消费者接管分区后，从上一次的偏移量开始消费
    @Override
    public void onPartitionsAssigned(Collection<TopicPartition> collection) {
        // 线程id
        final String threadId = Thread.currentThread().getId() + "";
        System.out.println(threadId + "|-再均衡完成，onPartitionsAssigned参数值为：" + collection);
        System.out.println("分区偏移量表中：" + partitionOffsetMap);

        for (TopicPartition topicPartition : collection) {
            System.out.println(threadId + "-topicPartition" + topicPartition);
            // 取得接管分区之前的偏移量
            Long offset = partitionOffsetMap.get(topicPartition);
            if (offset == null) continue;
            consumer.seek(topicPartition, partitionOffsetMap.get(topicPartition));
        }
    }

    @Override
    public void onPartitionsLost(Collection<TopicPartition> partitions) {
        ConsumerRebalanceListener.super.onPartitionsLost(partitions);
    }
}
